Replace websockets with zmq #9173

xjules · 2024-11-11T14:32:42Z

Issue

POC: testing feasibility of zmq.
Refactor tests

Approach
Implementing router-dealer pattern with custom acknowledgments with zmq

Server / router (ensemble evaluator) gets messages and keeps track of connected clients wherein in additional it sends ack message to dealaers to confirm that the message has been received .
Dealer(s) = clients and dispatchers are sending messages to the router via:

send_multipart([b"", message.encode("utf-8")])

dealers always wait for acknowledgment from the evaluator
removing websockets, no more wait_for_evaluator
Settup encryption with curve
each dealer (client, dispatcher) will get a unique name
Make sure to check cancellation error when sending event from client
Monitor is an advanced version Client
Use TCP protocol only when using LSF, SLURM or TORQUE queues
PR title captures the intent of the changes, and is fitting for release notes.
Added appropriate release note label
Commit history is consistent and clean, in line with the contribution guidelines.
Make sure unit tests pass locally after every commit (git rebase -i main --exec 'pytest tests/ert/unit_tests -n logical -m "not integration_test"')

When applicable

When there are user facing changes: Updated documentation
New behavior or changes to existing untested code: Ensured that unit tests are added (See Ground Rules).
Large PR: Prepare changes in small commits for more convenient review
Bug fix: Add regression test for the bug
Bug fix: Create Backport PR to latest release

src/_ert/forward_model_runner/client.py

src/ert/ensemble_evaluator/evaluator.py

codecov-commenter · 2024-12-10T15:14:02Z

Codecov Report

Attention: Patch coverage is 88.84892% with 31 lines in your changes missing coverage. Please review.

Project coverage is 91.71%. Comparing base (2a4b6be) to head (22b86ee).
Report is 23 commits behind head on main.

Files with missing lines	Patch %	Lines
src/_ert/forward_model_runner/client.py	85.39%	13 Missing ⚠️
src/ert/ensemble_evaluator/evaluator.py	87.12%	13 Missing ⚠️
src/_ert/forward_model_runner/reporting/event.py	91.42%	3 Missing ⚠️
src/ert/ensemble_evaluator/_ensemble.py	75.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9173      +/-   ##
==========================================
- Coverage   91.86%   91.71%   -0.15%     
==========================================
  Files         433      430       -3     
  Lines       26780    26684      -96     
==========================================
- Hits        24602    24474     -128     
- Misses       2178     2210      +32

Flag	Coverage Δ
cli-tests	`39.45% <65.56%> (-0.30%)`	⬇️
everest-models-test	`34.19% <74.82%> (-0.38%)`	⬇️
gui-tests	`71.75% <78.02%> (-0.37%)`	⬇️
integration-test	`38.34% <76.61%> (+1.13%)`	⬆️
performance-tests	`51.74% <78.75%> (-0.21%)`	⬇️
test	`39.63% <76.61%> (-1.05%)`	⬇️
unit-tests	`74.03% <86.33%> (-0.14%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

jonathan-eq · 2024-12-11T05:29:06Z

All tests passing 😻

jonathan-eq · 2024-12-12T09:17:43Z

src/_ert/forward_model_runner/client.py


    def __enter__(self) -> Self:
+        self.loop.run_until_complete(self.__aenter__())


Maybe we can remove all of the synced wrappers of the async methods?

Maybe. I think the synced version if now only used in tests.

~~I have made some exploration and the issue that the reporter tests are still fully synced and thus we need to preserve both (async and sync) versions alive.~~ Actually it's only the zmq socket that needs to be synced. Yes, this can go away!

jonathan-eq · 2024-12-12T09:20:50Z

src/_ert/forward_model_runner/client.py

    async def __aexit__(
        self, exc_type: Any, exc_value: Any, exc_traceback: Any


Do we need this in the method signature?
I see that we decorate the equivalent method in EnsembleEvaluator with contextlib.asynccontextmanager. Maybe we should do the same here if it is not used.

src/_ert/forward_model_runner/client.py

jonathan-eq · 2024-12-12T09:25:13Z

src/_ert/forward_model_runner/client.py

+    async def _term_receiver_task(self) -> None:
+        if self._receiver_task and not self._receiver_task.done():
+            self._receiver_task.cancel()
+            await asyncio.gather(self._receiver_task, return_exceptions=True)


Do we have to gather this cancelled task, or can we await it (and suppress asyncio.cancellationerror if it is raised)?

Since we capture asyncio.CancelledError we can just do await. :+1

Actually the CancelledError is not catched there, so I would keep it.

jonathan-eq · 2024-12-12T09:27:24Z

src/_ert/forward_model_runner/client.py

        return self

+    def term(self) -> None:


I think we should rearrange the order of the methods, it is a bit difficult to read

Can do 👍

src/_ert/forward_model_runner/client.py

jonathan-eq · 2024-12-13T07:42:02Z

src/_ert/forward_model_runner/reporting/event.py

+                while True:
+                    try:
+                        if self._done.is_set() and start_time is None:
+                            start_time = time.time()


What is start_time? Is that the start of closing the event_publisher?

Yes, that's the start_time of closing event_publisher. Maybe I can rename it start_closing_time?

Yes, that sounds good

src/_ert/forward_model_runner/reporting/event.py

src/ert/ensemble_evaluator/config.py

jonathan-eq · 2024-12-13T08:01:55Z

src/ert/ensemble_evaluator/evaluator.py

@@ -82,14 +65,20 @@ def __init__(self, ensemble: Ensemble, config: EvaluatorServerConfig):
        self._max_batch_size: int = 500
        self._batching_interval: float = 2.0
        self._complete_batch: asyncio.Event = asyncio.Event()
+        self._server_started: asyncio.Event = asyncio.Event()
+        self._clients_connected: set[bytes] = set()


I think self._connected_clients would be a better name. clients_connected sounds more like a boolean

Also, what is client? Is it the Client we use for sending/receiving, or is it referring to the monitor for UI?

Maybe we can rename it to monitors

Hm. I was trying to preserve the naming conventions we had before (which was clients), but maybe having it explicit as Monitor would make more sense. Then I would change handle_client -> handle_monitor as well.

Should we rename it now that we are already giving it an overhaul, or is that a separate PR? If not, let's change it to connected_clients

I would do it in another PR to keep the semantics as it was before.

src/ert/ensemble_evaluator/evaluator.py

jonathan-eq · 2024-12-13T08:26:50Z

src/ert/ensemble_evaluator/evaluator.py

+                self._router_socket.close()
+                zmq_context.destroy()


Suggested change

self._router_socket.close()

zmq_context.destroy()

zmq_context.destroy()

The destroy method should close all associated sockets automatically

I'll leave it for the time being as to be explicit.

src/ert/ensemble_evaluator/monitor.py

jonathan-eq · 2024-12-13T08:35:50Z

tests/ert/conftest.py

+            if sys.platform != "linux":
+                kwargs["use_ipc_protocol"] = True


What about the tests running on the GH Actions ubuntu runner?

This is meant only for MacOS GH Action runners.

Will we not use ipc on the ubuntu runners? Maybe that can speed up some tests

Yes, but we should test tcp sockets too. Hm, let's see.

src/ert/ensemble_evaluator/evaluator_connection_info.py

codspeed-hq · 2024-12-16T07:59:42Z

CodSpeed Performance Report

Merging #9173 will not alter performance

_{Comparing xjules:test_zmq (22b86ee) with main (1bd4fe1)}

Summary

✅ 24 untouched benchmarks

jonathan-eq · 2024-12-16T08:16:47Z

@xjules Seems codspeed is very impressed by your changes 🤣

jonathan-eq · 2024-12-13T09:33:51Z

src/_ert/forward_model_runner/client.py

+    async def process_message(self, msg: str) -> None:
+        pass


Can we throw on a @abstractmethod?
It is defined in monitor.py which inherits from Client

src/ert/ensemble_evaluator/evaluator.py

jonathan-eq · 2024-12-16T07:33:18Z

tests/ert/utils.py

+    def signal(self, value):
+        self.value = value


Suggested change

def signal(self, value):

self.value = value

def signal(self, should_stop: bool):

self.should_stop = should_stop

the signal(value) has more meaning than just "stop". Will put a comment there.

I've added a comment:

"""Mock ZMQ server for testing signal = 0: normal operation signal = 1: don't send ACK and don't receive messages signal = 2: don't send ACK, but receive messages """

jonathan-eq · 2024-12-16T14:04:39Z

src/_ert/forward_model_runner/reporting/event.py

+                while True:
+                    try:
+                        if self._done.is_set() and start_time is None:
+                            start_time = time.time()


Yes, that sounds good

src/ert/ensemble_evaluator/evaluator.py

jonathan-eq · 2024-12-19T09:39:50Z

tests/ert/conftest.py

+            if sys.platform != "linux":
+                kwargs["use_ipc_protocol"] = True


Will we not use ipc on the ubuntu runners? Maybe that can speed up some tests

tests/ert/unit_tests/ensemble_evaluator/test_monitor.py

tests/ert/unit_tests/shared/test_port_handler.py

tests/ert/utils.py

jonathan-eq

Very good, let's start the new year without websockets 🥇

- dealers always wait for acknowledgment from the evaluator - removing websockets and wait_for_evaluator - Settup encryption with curve - each dealer (client, dispatcher) will get a unique name - Monitor is an advanced version Client - _server_started.wait() is to signal that zmq router socket is bound - Use TCP protocol only when using LSF, SLURM or TORQUE queues -- Use ipc_protocol when using LOCAL driver - Remove certificate - Remove synced _send from Client - Remove cert generator - Remove ClientConnectionClosedOK - Add test for new connection while closing down evaluator - Add test for handle dispatcher and dispatcher messages in evaluator - Add tests for ipc and tcp ee config - Add test for clear connect and disconnect of Monitor - Set a a correct protocol for everestserver

xjules self-assigned this Nov 11, 2024

xjules force-pushed the test_zmq branch 3 times, most recently from feac786 to 5acae39 Compare November 19, 2024 14:58

xjules marked this pull request as ready for review November 19, 2024 14:59

jonathan-eq reviewed Nov 20, 2024

View reviewed changes

src/_ert/forward_model_runner/client.py Show resolved Hide resolved

jonathan-eq reviewed Nov 20, 2024

View reviewed changes

src/ert/ensemble_evaluator/evaluator.py Outdated Show resolved Hide resolved

xjules force-pushed the test_zmq branch 5 times, most recently from e553dbc to 016b43d Compare November 21, 2024 13:25

xjules force-pushed the test_zmq branch 4 times, most recently from ca1e52c to 0884d51 Compare December 6, 2024 13:09

xjules force-pushed the test_zmq branch 4 times, most recently from e3d6fa0 to 8b9a9ab Compare December 10, 2024 14:10

xjules force-pushed the test_zmq branch from 43a018a to 5787811 Compare December 11, 2024 00:01

jonathan-eq reviewed Dec 12, 2024

View reviewed changes

src/_ert/forward_model_runner/client.py Show resolved Hide resolved

jonathan-eq reviewed Dec 12, 2024

View reviewed changes

src/_ert/forward_model_runner/client.py Outdated Show resolved Hide resolved

jonathan-eq reviewed Dec 12, 2024

View reviewed changes

src/_ert/forward_model_runner/client.py Show resolved Hide resolved